import uuid
from typing import List, Dict, Optional, Any
from neo4j import GraphDatabase, Driver, Record
import logging

from neo4j.graph import Node

from config import *
from models.data_models import DialogNode

ROOT_NODE_NAMESPACE = uuid.UUID('f81d4fae-7dec-11d0-a765-00a0c91e6bf6')


def _map_record_to_dialog_node(record: Record) -> Optional[DialogNode]:
    if not record:
        return None
    node: Node = record.get('d')
    if not node:
        return None

    return DialogNode(
        id=uuid.UUID(node.get("id")),
        content=node.get("content", ""),
        depth=int(node.get("depth", 0)),
        embedding=node.get("embedding", [])
    )


class Neo4jManager:
    def __init__(self):
        try:
            self._driver: Driver = GraphDatabase.driver(
                NEO4J_URI,
                auth=(NEO4J_USERNAME, NEO4J_PASSWORD)
            )
            logging.info("Neo4j driver initialized successfully.")
            self.verify_connection()
        except Exception as e:
            logging.error(f"Failed to initialize Neo4j driver: {e}")
            self._driver = None

    def get_session_max_depth(self, session_id: uuid.UUID) -> int:
        if not self._driver:
            raise ConnectionError("Neo4j driver is not available.")
        def _tx_get_max_depth(tx, session_id_str):
            query = """
            MATCH (s:Session {id: $session_id})-[:HAS_ROOT]->(root)
            MATCH p = (root)-[:HAS_CHILD*0..]->(leaf)
            WHERE NOT (leaf)-[:HAS_CHILD]->()
            RETURN max(length(p)) AS maxDepth
            """
            result = tx.run(query, session_id=session_id_str)
            record = result.single()
            return record["maxDepth"] if record and record["maxDepth"] is not None else 0
        with self._driver.session() as session:
            return session.execute_read(_tx_get_max_depth, str(session_id))

    def close(self):
        if self._driver:
            self._driver.close()
            logging.info("Neo4j driver closed.")

    def delete_node(self, dialog_id: uuid.UUID):
        if not self._driver:
            raise ConnectionError("Neo4j driver is not available.")

        def _tx_delete_node(tx, dialog_id_str):
            query = "MATCH (d:Dialog {id: $dialog_id}) DETACH DELETE d"
            tx.run(query, dialog_id=dialog_id_str)

        with self._driver.session() as session:
            session.execute_write(_tx_delete_node, str(dialog_id))
            logging.info(f"Deleted dialog node: {dialog_id}")

    def verify_connection(self):
        if not self._driver:
            raise ConnectionError("Neo4j driver is not available.")
        try:
            self._driver.verify_connectivity()
            logging.info("Successfully connected to Neo4j database.")
        except Exception as e:
            logging.error(f"Failed to connect to Neo4j: {e}")
            raise ConnectionError(f"Could not connect to Neo4j: {e}") from e

    def find_or_create_root_dialog(self, session_id: uuid.UUID) -> DialogNode:
        if not self._driver:
            raise ConnectionError("Neo4j driver is not available.")

        root_id = uuid.uuid5(ROOT_NODE_NAMESPACE, str(session_id))

        def _tx_find_or_create_root(tx, session_id_str, root_id_str):
            query = """
            MERGE (s:Session {id: $session_id})
            MERGE (s)-[:HAS_ROOT]->(d:Dialog {id: $root_id})
            ON CREATE SET 
                d.content = 'session root',
                d.depth = 0,
                d.embedding = [],
                d.id = $root_id
            RETURN d
            """
            result = tx.run(query, session_id=session_id_str, root_id=root_id_str)
            return result.single()

        with self._driver.session() as session:
            record = session.execute_write(_tx_find_or_create_root, str(session_id), str(root_id))
            return _map_record_to_dialog_node(record)

    def load_dialog(self, dialog_id: uuid.UUID) -> Optional[DialogNode]:
        if not self._driver:
            raise ConnectionError("Neo4j driver is not available.")

        def _tx_load_dialog(tx, dialog_id_str):
            query = "MATCH (d:Dialog {id: $dialog_id}) RETURN d"
            result = tx.run(query, dialog_id=dialog_id_str)
            return result.single()

        with self._driver.session() as session:
            record = session.execute_read(_tx_load_dialog, str(dialog_id))
            return _map_record_to_dialog_node(record)

    def load_children_dialog_ids(self, parent_id: uuid.UUID) -> List[uuid.UUID]:
        if not self._driver:
            raise ConnectionError("Neo4j driver is not available.")

        def _tx_load_children_ids(tx, parent_id_str):
            query = """
            MATCH (p:Dialog {id: $parent_id})-[:HAS_CHILD]->(c:Dialog)
            RETURN c.id AS child_id
            """
            results = tx.run(query, parent_id=parent_id_str)
            
            return [uuid.UUID(record["child_id"]) for record in results]

        with self._driver.session() as session:
            return session.execute_read(_tx_load_children_ids, str(parent_id))


    def batch_update_dialogs(self, updates: list[dict]):
        if not self._driver or not updates:
            return

        def _tx_batch_update(tx, update_list):
            query = """
            UNWIND $updates AS update_item
            MATCH (d:Dialog {id: update_item.id})
            SET d.content = update_item.summary, d.embedding = update_item.embedding
            """
            tx.run(query, updates=update_list)

        with self._driver.session() as session:
            session.execute_write(_tx_batch_update, updates)
        logging.info(f"Batch updated {len(updates)} dialog nodes in Neo4j.")

    def insert_dialog(self, parent_id: uuid.UUID, content: str, depth: int, embedding: List[float]) -> DialogNode:
        if not self._driver:
            raise ConnectionError("Neo4j driver is not available.")

        new_id = uuid.uuid4()

        def _tx_insert_dialog(tx, parent_id_str, dialog_props):
            query = """
            MATCH (p:Dialog {id: $parent_id})
            CREATE (d:Dialog)
            SET d = $props
            CREATE (p)-[:HAS_CHILD]->(d)
            RETURN d
            """
            result = tx.run(query, parent_id=parent_id_str, props=dialog_props)
            return result.single()

        dialog_properties = {
            "id": str(new_id),
            "content": content,
            "depth": depth,
            "embedding": embedding
        }

        with self._driver.session() as session:
            record = session.execute_write(_tx_insert_dialog, str(parent_id), dialog_properties)
            return _map_record_to_dialog_node(record)

    def update_dialog_summary_and_embedding(self, dialog_id: uuid.UUID, summary: str, embedding: List[float]):
        if not self._driver:
            raise ConnectionError("Neo4j driver is not available.")

        def _tx_update_dialog(tx, dialog_id_str, summary_str, embedding_list):
            query = """
            MATCH (d:Dialog {id: $dialog_id})
            SET d.content = $summary, d.embedding = $embedding
            """
            tx.run(query, dialog_id=dialog_id_str, summary=summary_str, embedding=embedding_list)

        with self._driver.session() as session:
            session.execute_write(_tx_update_dialog, str(dialog_id), summary, embedding)
            logging.info(f"Updated summary for dialog node: {dialog_id}")

    def load_session_vectors_for_faiss(self, session_id: uuid.UUID) -> List[Dict[str, Any]]:
        if not self._driver:
            raise ConnectionError("Neo4j driver is not available.")

        def _tx_load_vectors(tx, session_id_str):
            query = """
            MATCH (:Session {id: $session_id})-[:HAS_ROOT]->(root:Dialog)
            MATCH (root)-[:HAS_CHILD*0..]->(node:Dialog)
            WHERE node.embedding IS NOT NULL AND size(node.embedding) > 0
            RETURN node.id AS id, node.embedding AS embedding
            """
            results = tx.run(query, session_id=session_id_str)
            return [
                {"id": uuid.UUID(record["id"]), "embedding": record["embedding"]}
                for record in results
            ]

        with self._driver.session() as session:
            return session.execute_read(_tx_load_vectors, str(session_id))